{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cart-Pole: Linearization and LQR Balncing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib.\n",
    "from IPython import get_ipython\n",
    "if get_ipython() is not None: get_ipython().run_line_magic(\"matplotlib\", \"inline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "# underactuated imports\n",
    "from underactuated import FindResource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "In this problem you will work on the cart-pole system described in [Section 3.2 of the textbook](http://underactuated.csail.mit.edu/acrobot.html#cart_pole).\n",
    "You will be asked to write down its dynamics in state-space form, linearized it, and then analyze the linearization error.\n",
    "At the end we will wire up an LQR controller and we will simulate the cart-pole in a series of balancing tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics in State-Space Form\n",
    "Consider the cart-pole system described in [Section 3.2 of the textbook](http://underactuated.csail.mit.edu/acrobot.html#cart_pole).\n",
    "For the sake of simplicity, in this notebook we fix the following numeric values for its parameters:\n",
    "- mass of the cart $m_{\\text{c}}=1$,\n",
    "- mass of the pole $m_{\\text{p}}=1$,\n",
    "- length of the pole $l=1$,\n",
    "- gravity acceleration $g=9.81$.\n",
    "\n",
    "**Important:** Do not modify/round these parameters, otherwise the autograding code will raise an error.\n",
    "\n",
    "Using the same convention of the book, we describe the state as $$\\mathbf{x} = [x, \\theta, \\dot{x}, \\dot{\\theta}]^T,$$ and we let the force on the cart be the control input $\\mathbf{u} = [f]$.\n",
    "Use [equations (16) and (17) from the textbook](http://underactuated.csail.mit.edu/acrobot.html#cart_pole) to derive the state-space model of the cart-pole: $$\\dot{\\mathbf{x}} = f(\\mathbf{x}, \\mathbf{u}).$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the state x (4d array)\n",
    "# and the input u (1d array) returns the right\n",
    "# hand side of the state space dynamics f(x,u)\n",
    "# (remember that we fixes the cart-pole parameters\n",
    "# to the values above!)\n",
    "def f(x, u):\n",
    "    \n",
    "    # shortcuts for the cosine and the sine of theta\n",
    "    # they might be handy\n",
    "    c = np.cos(x[1])\n",
    "    s = np.sin(x[1])\n",
    "    \n",
    "    # gravity acceleration\n",
    "    g = 9.81 # do not change\n",
    "        \n",
    "    # fill the following matrix\n",
    "    # (sorry for the one-base counting!)\n",
    "    f1 = 0 # modify here\n",
    "    f2 = 0 # modify here\n",
    "    f3 = 0 # modify here\n",
    "    f4 = 0 # modify here\n",
    "    \n",
    "    return np.array([f1, f2, f3, f4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics Linearized Around the Unstable Equilibrium\n",
    "We now approximate the nonlinear dynamics with a linear one.\n",
    "This will allow us to basic linear control to stabilize (locally) the cart-pole with the pole in the vertical configuration.\n",
    "\n",
    "We consider the unstable equilibrium state $$\\mathbf{x}^* = [0, \\pi, 0, 0]^T,$$ with the related equilibrium control input $$\\mathbf{u}^* = [0].$$\n",
    "As in [Section 3.4.1](http://underactuated.csail.mit.edu/acrobot.html#section4), we want to derive a linear model in the from\n",
    "$$\\dot{\\bar{\\mathbf{x}}} = A_{\\text{lin}} \\mathbf{\\bar{x}} + B_{\\text{lin}} \\mathbf{\\bar{u}},$$\n",
    "where $\\mathbf{\\bar{x}} = \\mathbf{x}-\\mathbf{x}^*$ and $\\mathbf{\\bar{u}} = \\mathbf{u} -\\mathbf{u}^*$.\n",
    "\n",
    "Follow the recipe described in [Section 3.4.1 of the textbook](http://underactuated.csail.mit.edu/acrobot.html#section4) to derive the linearization matrices $A_{\\text{lin}}$ and $B_{\\text{lin}}$, and implement them in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that returns the A_lin matrix\n",
    "def get_A_lin():\n",
    "    g = 9.81 # do not change\n",
    "    # fill the matrix below\n",
    "    A = np.array([\n",
    "        [0, 0, 0, 0], # modify here\n",
    "        [0, 0, 0, 0], # modify here\n",
    "        [0, 0, 0, 0], # modify here\n",
    "        [0, 0, 0, 0]  # modify here\n",
    "    ])\n",
    "    return A\n",
    "    \n",
    "# function that returns the B_lin matrix\n",
    "def get_B_lin():\n",
    "    # fill the matrix below\n",
    "    B = np.array([\n",
    "        [0], # modify here\n",
    "        [0], # modify here\n",
    "        [0], # modify here\n",
    "        [0]  # modify here\n",
    "    ])\n",
    "    return B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linearization Error\n",
    "The linear model we have built above is very accurate accurate in the vicinity of the equilibrium point, but can lead to very bad predictions if our state is far away from the equilibrium.\n",
    "\n",
    "The following function, for a given state $\\mathbf{x}$ and control $\\mathbf{u}$, evaluates the linearization error:\n",
    "$$\n",
    "e(\\mathbf{x}, \\mathbf{u})\n",
    "=\n",
    "\\| f(\\mathbf{x}, \\mathbf{u}) - f_{\\text{lin}}(\\mathbf{x}, \\mathbf{u}) \\|,$$\n",
    "where we defined $f_{\\text{lin}}(\\mathbf{x}, \\mathbf{u}) = A_{\\text{lin}} \\mathbf{\\bar{x}} + B_{\\text{lin}} \\mathbf{\\bar{u}}.$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_lin(x,u):\n",
    "    \n",
    "    # equilibrium point\n",
    "    x_star = np.array([0, np.pi, 0, 0])\n",
    "    u_star = np.array([0])\n",
    "    \n",
    "    # linearized dynamics\n",
    "    x_bar = x - x_star\n",
    "    u_bar = u - u_star\n",
    "    A = get_A_lin()\n",
    "    B = get_B_lin()\n",
    "    \n",
    "    return A.dot(x_bar) + B.dot(u_bar)\n",
    "    \n",
    "def linearization_error(x, u):\n",
    "    return np.linalg.norm(f(x,u) - f_lin(x,u))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the function above to evaluate the error $e(\\mathbf{x}, \\mathbf{u})$ in the following 6 conditions:\n",
    "- $\\mathbf{x} = [0, 0.99 \\pi, 0, 0]^T$ and $\\mathbf{u} = [0]$,\n",
    "- $\\mathbf{x} = [0, 0.9 \\pi, 0, 0]^T$ and $\\mathbf{u} = [-10]$,\n",
    "- $\\mathbf{x} = [0, 0.85 \\pi, 0, 0]^T$ and $\\mathbf{u} = [0]$,\n",
    "- $\\mathbf{x} = [0, 0.5 \\pi, 0, 0]^T$ and $\\mathbf{u} = [0]$,\n",
    "- $\\mathbf{x} = [0, 0, 0, 0]^T$ and $\\mathbf{u} = [0]$,\n",
    "- $\\mathbf{x} = [1, \\pi, 0, 0]^T$ and $\\mathbf{u} = [10]$,\n",
    "\n",
    "**Attention 1:** For the number $\\pi$ use `np.pi`! **Do not** truncate it like $3.14$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fill these states with the ones given above\n",
    "x_list = [\n",
    "    np.array([0, np.pi, 0, 0]), # modify here\n",
    "    np.array([0, np.pi, 0, 0]), # modify here\n",
    "    np.array([0, np.pi, 0, 0]), # modify here\n",
    "    np.array([0, np.pi, 0, 0]), # modify here\n",
    "    np.array([0, np.pi, 0, 0]), # modify here\n",
    "    np.array([0, np.pi, 0, 0])  # modify here\n",
    "]\n",
    "\n",
    "# fill these inputs with the ones given above\n",
    "u_list = [\n",
    "    np.array([0]), # modify here\n",
    "    np.array([0]), # modify here\n",
    "    np.array([0]), # modify here\n",
    "    np.array([0]), # modify here\n",
    "    np.array([0]), # modify here\n",
    "    np.array([0])  # modify here\n",
    "]\n",
    "\n",
    "# compute linearization errors for all the points above\n",
    "errors = [linearization_error(x_list[i], u_list[i]) for i in range(6)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we compare these linearization errors with the norm of $\\dot{\\mathbf{x}}$, i.e.,\n",
    "$$\n",
    "\\| f(\\mathbf{x}, \\mathbf{u})\\|.$$\n",
    "Clearly, the smaller is the linearization error with respect to this value, the better is our linear model.\n",
    "\n",
    "Spend the time you need to convince yourself about this result.\n",
    "Try to answer the following questions:\n",
    "- Is our linear approximation valid for all the points we tested?\n",
    "- Do we expect a linear controller to do a decent job when $\\theta = \\pi/2$?\n",
    "- When $\\theta$ is different from zero, does the linearization error depend on $\\mathbf{u}$?\n",
    "- Why is the error from the second case bigger than the one from the third, even if the second $\\theta$ is closer to $\\pi$ than the third?\n",
    "- What about the position $x$ of the cart? Should it affect the linearization error? If no, why not?\n",
    "\n",
    "(Questions not graded, do not submit.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, e in enumerate(errors):\n",
    "    print(f'State = {np.around(x_list[i], decimals=3)}^T')\n",
    "    print(f'Input = {np.around(u_list[i], decimals=3)}')\n",
    "    print('Linearization error = {:.3f}'.format(e))\n",
    "    print('Norm of f(x,u) = {:.3f}\\n'.format(np.linalg.norm(f(x_list[i], u_list[i]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Balancing with LQR Controller\n",
    "We finally move to the design of the LQR controller.\n",
    "Drake handles all the linearization process very transparently: no need to get your hands dirty with all the linearization issues we have discussed above!\n",
    "But it was worth to do it by hand at least once...\n",
    "\n",
    "Drake can design an LQR controller directly on the nonlinear system obtained by parsing the `.urdf` file!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pydrake imports\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, DiagramBuilder,\n",
    "                         LinearQuadraticRegulator, Parser,\n",
    "                         PlanarSceneGraphVisualizer, Simulator, Linearize)\n",
    "\n",
    "# underactuated imports\n",
    "from underactuated import FindResource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we set a couple of numeric parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unstable equilibrium point\n",
    "x_star = [0, np.pi, 0, 0]\n",
    "\n",
    "# weight matrices for the lqr controller\n",
    "Q = np.eye(4)\n",
    "R = np.eye(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we construct the block diagram with the cart-pole in closed loop with the LQR controller."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start construction site of our block diagram\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "# instantiate the cart-pole and the scene graph\n",
    "cartpole, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "urdf_path = FindResource('models/undamped_cartpole.urdf')\n",
    "Parser(cartpole).AddModelFromFile(urdf_path)\n",
    "cartpole.Finalize()\n",
    "\n",
    "# set the operating point (vertical unstable equilibrium)\n",
    "context = cartpole.CreateDefaultContext()\n",
    "context.get_mutable_continuous_state_vector().SetFromVector(x_star)\n",
    "\n",
    "# fix the input port to zero and get its index for the lqr function\n",
    "cartpole.get_actuation_input_port().FixValue(context, [0])\n",
    "input_i = cartpole.get_actuation_input_port().get_index()\n",
    "\n",
    "# synthesize lqr controller directly from\n",
    "# the nonlinear system and the operating point\n",
    "lqr = LinearQuadraticRegulator(cartpole, context, Q, R, input_port_index=input_i)\n",
    "lqr = builder.AddSystem(lqr)\n",
    "\n",
    "# the following two lines are not needed here...\n",
    "output_i = cartpole.get_state_output_port().get_index()\n",
    "cartpole_lin = Linearize(cartpole, context, input_port_index=input_i, output_port_index=output_i)\n",
    "\n",
    "# wire cart-pole and lqr\n",
    "builder.Connect(cartpole.get_state_output_port(), lqr.get_input_port(0))\n",
    "builder.Connect(lqr.get_output_port(0), cartpole.get_actuation_input_port())\n",
    "\n",
    "# add a visualizer and wire it\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-3., 3.], ylim=[-1.2, 1.2], show=False)\n",
    ")\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(), visualizer.get_input_port(0))\n",
    "\n",
    "# finish building the block diagram\n",
    "diagram = builder.Build()\n",
    "\n",
    "# instantiate a simulator\n",
    "simulator = Simulator(diagram)\n",
    "simulator.set_publish_every_time_step(False) # makes sim faster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell contains a function that simulates the closed-loop system and produces a video of the sim."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the cart-pole initial state\n",
    "# and the simulation time, simulates the system\n",
    "# and produces a video\n",
    "def simulate_and_animate(x0, sim_time=5):\n",
    "    \n",
    "    # start recording the video for the animation of the simulation\n",
    "    visualizer.start_recording()\n",
    "    \n",
    "    # reset initial time and state\n",
    "    context = simulator.get_mutable_context()\n",
    "    context.SetTime(0.)\n",
    "    context.SetContinuousState(x0)\n",
    "    \n",
    "    # run sim\n",
    "    simulator.Initialize()\n",
    "    simulator.AdvanceTo(sim_time)\n",
    "    \n",
    "    # stop video\n",
    "    visualizer.stop_recording()\n",
    "    \n",
    "    # construct animation\n",
    "    ani = visualizer.get_recording_as_animation()\n",
    "    \n",
    "    # display animation below the cell\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "    \n",
    "    # reset to empty video\n",
    "    visualizer.reset_recording()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we just run the function we just wrote for all the initial states we analyzed in this notebook, and we look at the result!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate and animate the cart\n",
    "for x in x_list:\n",
    "    simulate_and_animate(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Was your intuition from the previous analysis correct?\n",
    "Out of the 6 initial states we considered, which are the states from which the LQR controller is able to recover? (Questions not graded, do not submit.)\n",
    "\n",
    "In the next cell write (in base zero) the indices of the states from which the system is able to recover (autograded)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_recovers_from_states = [] # modify here\n",
    "print('System recovers from states:')\n",
    "for i in system_recovers_from_states:\n",
    "    print(np.around(x_list[i], decimals=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Final Note\n",
    "In the middle of the construction of the block diagram above, we have hidden the system `cartpole_lin`.\n",
    "It has been defined using `Linearize`.\n",
    "This is the function that `LinearQuadraticRegulator` uses to linearize the plant before solving the Riccati equation. \n",
    "Feel free to use the methods `cartpole_lin.A()` and `cartpole_lin.B()` to double check your answer above!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "You can check your work by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.acrobot.cartpole_balancing.test_cartpole_balancing import TestCartPoleBalancing\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestCartPoleBalancing], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}